#include "Database.h"

using namespace std;
using namespace clUtil;

static stringstream exceptionString("");
static const unsigned int kMaxPeptideLength = 63;

static char aminoAcids[] = {'R', 'H', 'K', 'D', 'E', 'S', 'T', 'N', 'Q', 'C', 
                            'U', 'G', 'P', 'A', 'I', 'L', 'M', 'F', 'W', 'Y', 
                            'V'};

//Makes alpha characters uppercase. If not a lowercase character, just return it
static inline char toUpper(char input)
{
  return input >= 'a' && input <= 'z' ? input - ('a' - 'A') : input;
}

static bool isValidCode(char curCharacter)
{
  //Seriously abuse C's fallthrough of case statements...
  switch(curCharacter)
  {
    case 'A'://Alanine
    case 'B'://Aspartic acid
    case 'C'://Cysteine
    case 'D'://Aspartic acid
    case 'E'://Glutamic acid
    case 'F'://Phenylalanine
    case 'G'://Glycine
    case 'H'://Histidine
    case 'I'://Isoleucine
    case 'K'://Lysine
    case 'L'://Leucine
    case 'M'://Methionine
    case 'N'://Asparagine
    case 'P'://Proline
    case 'Q'://Glutamine
    case 'R'://Arginine
    case 'S'://Serine
    case 'T'://Threonine
    case 'V'://Valine
    case 'W'://Tryptophan
    case 'Y'://Tyrosine
    case 'Z'://Glutamic acid
    case 'X'://Special: ANY
    case '*'://Special: TRANSLATION STOP
      return true;
    case '-'://Special: GAP
      fprintf(stderr, "Error: specmaster does not support gaps.\n");
    default:
      return false;
  }
}

static void appendSequence(string& theString, DatabaseEntry& entry)
{
  for(size_t i = 0; i < theString.size(); i++)
  {
    char curChar = toUpper(theString[i]);

    if(isValidCode(curChar) == true)
    {
      if(curChar == '*')
      {
        break;
      }
      if(curChar == 'X')
      {
        entry.xCount++;
      }
      entry.sequence += curChar;
    }
    else
    {
      exceptionString.str("");
      exceptionString << "Invalid code '" << curChar << "' column " << i;
      throw runtime_error(exceptionString.str());
    }
  }
}

bool Database::processEntry(DatabaseEntry& entry, unsigned int recursionLevel)
{
  bool bufferFull = false;

  //If entry has an X, replace the first X with every possible amino acid.
  //Recursively process each permutation to remove all the Xs.
  //Print a warning because this is bad and grows exponentially.
  if(entry.xCount > 0)
  {
    if(recursionLevel == 0)
    {
      cout << "WARNING! Protein with wildcard ('X') found:\n"
           << entry.comment
           << "\n\tUse extreme caution as these will cause processing time of"
           << "\n\tthis amino acid to grow exponentially\n";
    }
    for(unsigned int curAcid = 0; curAcid < sizeof(aminoAcids); curAcid++)
    {
      DatabaseEntry newEntry;
      newEntry.xCount = entry.xCount - 1;
      newEntry.comment = entry.comment + aminoAcids[curAcid];
      newEntry.sequence = entry.sequence;

      //Find the x...
      for(unsigned int i = 0; i < entry.sequence.size(); i++)
      {
        if(newEntry.sequence[i] == 'X')
        {
          newEntry.sequence[i] = aminoAcids[curAcid];
          break;
        }
      }

      bufferFull = processEntry(newEntry, recursionLevel + 1);
    }

    
  }
  else
  {
    //Cleave the protein...
    vector<string> curPeptides;
    unsigned int newBufferSize;
    
    dissolveProtein(entry.sequence, curPeptides, mCleavageFunction);
    newBufferSize = (curPeptides.size() + mAllPeptides.size()) * 
      (kMaxPeptideLength + 1);
    
    if(newBufferSize > gBufferSize)
    {
      return true;
    }

    //For each peptide in the current protein
    for(size_t j = 0; j < curPeptides.size(); j++)
    {
      string& curPeptideStr = curPeptides[j];

      //If peptide is not in multimap, append it to main peptide list
      if(curPeptideStr.size() <= kMaxPeptideLength && 
         mPeptideMap.find(curPeptideStr) == mPeptideMap.end())
      {
        Peptide curPeptide(curPeptideStr);

        mAllPeptides.push_back(curPeptideStr);
      
        mPeptideMap.insert(pair<string,string>(curPeptideStr, entry.name));
      }
      else if(curPeptideStr.size() <= kMaxPeptideLength)
      {
        mPeptideMap.insert(pair<string,string>(curPeptideStr, entry.name));
      }
    }
  }

  return bufferFull;
}

Database::Database(const char* filename, bool (*cleavageFunction)(char, char)) :
  mCleavageFunction(cleavageFunction)
{
  mCurProtein = 0;
  ifstream databaseFile(filename);
  size_t numEntries = 0;
  size_t currentLineNum = 0;
  bool firstSequence = true;
  bool finishedHeader = false;
  DatabaseEntry curEntry;

  if(databaseFile.fail() == true)
  {
    throw runtime_error("Could not open database file.\n");
  }

  //Read parse the database file
  while(1)
  {
    string currentLine;

    //Read until end of file
    if(databaseFile.good() == false)
    {
      //Add the last entry to the database
      mAllProteins.push_back(curEntry);
      numEntries++;
      break;
    }

    getline(databaseFile, currentLine);
    
    //Stupid Windows and its carriage returns...
    if(currentLine[currentLine.size() - 1] == '\r')
    {
      currentLine.erase(currentLine.end() - 1);
    }

    currentLineNum++;

    if(currentLine[0] == '>' || currentLine[0] == ';')
    {
      
      //Handle reading the first database entry
      if(firstSequence == false && finishedHeader == true)
      {
        mAllProteins.push_back(curEntry);
        numEntries++;
      }

      //Continue previous entry
      if(finishedHeader == false)
      {
        curEntry.comment += currentLine;
      }
      //New entry
      else
      {
        curEntry.comment = currentLine.substr(1, currentLine.size());
        curEntry.sequence = "";
        curEntry.xCount = 0;
      }

      finishedHeader = false;
    }
    else
    {
      try
      {
        appendSequence(currentLine, curEntry);
        firstSequence = false;
        finishedHeader = true;
        curEntry.name = curEntry.comment.substr(0, curEntry.comment.find(" "));
      }
      catch(exception& err)
      {
        exceptionString.str("");
        exceptionString << __func__ 
                        << ":Error line " 
                        << currentLineNum 
                        << endl;
        exceptionString << "\t" << err.what();
        
        throw runtime_error(exceptionString.str());
      }
    }
  }

  databaseFile.close();
}

bool Database::copyToDevices()
{
  size_t numDevices = Device::GetDevices().size();
  size_t numEntries = 0;
  bool proteinsDone = false;

  mAllPeptides.clear();

  //Chop up the strings and make distractor set
  for(unsigned int i = mCurProtein;i<mAllProteins.size();i++)
  {
    string forwardProtein = (mAllProteins[i]).sequence;
    string reverseProtein = forwardProtein;

    if(processEntry(mAllProteins[i]))
    {
      mCurProtein = i;
      break;
    }
    numEntries++;
    //Reverse the protein string
    for(unsigned int j = 0; j < forwardProtein.size(); j++)
    {
      reverseProtein[forwardProtein.size() - 1 - j] = forwardProtein[j];
    }

    DatabaseEntry reverseEntry;
    
    reverseEntry.comment = "rev_" + (mAllProteins[i]).comment;
    reverseEntry.name = "rev_" + (mAllProteins[i]).name;
    reverseEntry.sequence = reverseProtein;
    
    if(processEntry(reverseEntry))
    {
      mCurProtein = i;
      break;
    }
    numEntries++;
    
    if(i == mAllProteins.size() -1)
    {
      proteinsDone = true;
      mCurProtein = i;
    }
  }

  cout <<"Processed " << mCurProtein + 1 << " out of " 
       << mAllProteins.size() << " proteins" << endl;
  
  //Sort all the peptides by their masses
  sort(mAllPeptides.begin(), mAllPeptides.end());
  
  //Allocate the peptide database at 1+63 bytes per entry. 1st byte is length. 
  //Assume no peptide is longer than 63 characters. This is what we pass
  //to OpenCL devices...
  mNumPeptides = mAllPeptides.size();

  //Allocate database and masses
  unique_ptr<char[]> peptides(new char[mNumPeptides * (kMaxPeptideLength + 1)]);
  unique_ptr<unsigned int[]> peptideMasses(new unsigned int[mNumPeptides]);

  //Give ownership to this class
  mPeptideDatabase = move(peptides);
  mPeptideMasses = move(peptideMasses);

  cout << mNumPeptides << " peptides in " << numEntries << " proteins" << endl;

  //Copy peptides into final database
  for(size_t i = 0; i < mNumPeptides; i++)
  {
    string& curString = mAllPeptides[i].mAcids;

    //Write string length to first byte
    mPeptideDatabase[(kMaxPeptideLength + 1)*i + 0] = 
      (unsigned char)curString.size();

    //Write string
    for(size_t j = 0; j < curString.size(); j++)
    {
      mPeptideDatabase[(kMaxPeptideLength + 1)*i + j + 1] = curString[j];
    }

    //Write NULL characters to rest of string
    for(size_t j = curString.size(); j < kMaxPeptideLength; j++)
    {
      mPeptideDatabase[(kMaxPeptideLength + 1)*i + j + 1] = 'B';
    }

    mPeptideMasses[i] = mAllPeptides[i].mMass;
  }

  mDevicePeptideDatabase.clear();
  mDevicePeptideMasses.clear();

  //Allocate and copy peptides to device
  for(unsigned int i = 0; i < numDevices; i++)
  {
    Device::SetCurrentDevice(i);

    //Allocate the database image
    unique_ptr<Image> tmpDatabase(new Image(4 * mNumPeptides, 
                                            CL_RGBA, 
                                            CL_FLOAT));
    tmpDatabase->put(mPeptideDatabase.get());

    //Allocate the mass buffer
    unique_ptr<Buffer> tmpMasses(new Buffer(sizeof(unsigned int) * 
                                            mNumPeptides));
    tmpMasses->put(mPeptideMasses.get());

    mDevicePeptideDatabase.push_back(move(tmpDatabase));
    mDevicePeptideMasses.push_back(move(tmpMasses));
  }

  return proteinsDone;
}

Peptide& Database::getPeptide(unsigned int num)
{
  return mAllPeptides[num];
}

void Database::printPeptides()
{
  for(unsigned int curPeptide = 0; curPeptide < mNumPeptides; curPeptide++)
  {
    char* curPeptidePtr = 
      &mPeptideDatabase[(kMaxPeptideLength + 1) * curPeptide];

    for(unsigned char curAA = 1; 
        curAA <= curPeptidePtr[0];
        curAA++)
    {
      printf("%c", curPeptidePtr[curAA]);
    }

    printf("\t%u\n", mPeptideMasses[curPeptide]);
  }
}

pair<multimap<string, string>::iterator, multimap<string, string>::iterator> 
Database::getProtein(string& peptide)
{
  return mPeptideMap.equal_range(peptide);
}
